{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Tensorflow 前端\n",
        "\n",
        "参考: [](https://xinetzone.github.io/tvm/docs/arch/frontend/tensorflow.html)\n",
        "\n",
        "```{note}\n",
        "请将 `tensorflow` 的 GPU 内存使用限制在必要的范围内，而不是使用所有可用的内存。您可以参考 [limiting_gpu_memory_growth](https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth) 了解如何进行操作。\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2023-06-09 14:12:16.956157: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
            "2023-06-09 14:12:17.005059: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
            "2023-06-09 14:12:17.006107: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
            "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
            "2023-06-09 14:12:17.794152: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
            "2023-06-09 14:12:19.241479: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1956] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
            "Skipping registering GPU devices...\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "try:\n",
        "    tf_compat_v1 = tf.compat.v1\n",
        "except (ImportError, AttributeError):\n",
        "    tf_compat_v1 = tf\n",
        "\n",
        "gpus = tf.config.list_physical_devices(\"GPU\")\n",
        "if gpus:\n",
        "    try:\n",
        "        for gpu in gpus:\n",
        "            tf.config.experimental.set_memory_growth(gpu, True)\n",
        "        print(\"tensorflow will use experimental.set_memory_growth(True)\")\n",
        "    except RuntimeError as e:\n",
        "        print(\"experimental.set_memory_growth option is not available: {}\".format(e))\n",
        "# Tensorflow 实用函数\n",
        "import tvm.relay.testing.tf as tf_testing"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 准备阶段"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# 模型相关文件的基础位置。\n",
        "repo_base = \"https://github.com/dmlc/web-data/raw/main/tensorflow/models/InceptionV1\"\n",
        "# 测试图片\n",
        "img_name = \"elephant-299.jpg\"\n",
        "image_url = f\"{repo_base}/{img_name}\"\n",
        "# 模型\n",
        "model_name = \"classify_image_graph_def-with_shapes.pb\"\n",
        "model_url = f\"{repo_base}/{model_name}\"\n",
        "# 图像标签映射\n",
        "map_proto = \"imagenet_2012_challenge_label_map_proto.pbtxt\"\n",
        "map_proto_url = f\"{repo_base}/{map_proto}\"\n",
        "# 标签的人类可读文本。\n",
        "label_map = \"imagenet_synset_to_human_label_map.txt\"\n",
        "label_map_url = f\"{repo_base}/{label_map}\""
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "下载如下文件："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib.download import download_testdata\n",
        "\n",
        "image_path = download_testdata(image_url, img_name, module=\"data\")\n",
        "model_path = download_testdata(model_url, model_name, module=[\"tf\", \"InceptionV1\"])\n",
        "map_proto_path = download_testdata(map_proto_url, map_proto, module=\"data\")\n",
        "label_path = download_testdata(label_map_url, label_map, module=\"data\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 在 tensorflow 上推理\n",
        "\n",
        "在 TensorFlow 上运行相应的模型。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_graph(model_path):\n",
        "    \"\"\"从保存的 GraphDef 文件创建图\"\"\"\n",
        "    # 从保存的 graph_def.pb 文件创建图。\n",
        "    with tf_compat_v1.gfile.GFile(model_path, \"rb\") as f:\n",
        "        graph_def = tf_compat_v1.GraphDef()\n",
        "        graph_def.ParseFromString(f.read())\n",
        "        # 将 graph_def 中的图导入到当前默认的图中。\n",
        "        tf.import_graph_def(graph_def, name=\"\") \n",
        "        # 对 `graph_def` 进行类型检查，可能进行规范化。\n",
        "        graph_def = tf_testing.ProcessGraphDefParam(graph_def)\n",
        "    return graph_def\n",
        "\n",
        "def read_image_tf(image_path):\n",
        "    if not tf_compat_v1.gfile.Exists(image_path):\n",
        "        tf.logging.fatal(\"File does not exist %s\", image_path)\n",
        "    with tf_compat_v1.gfile.GFile(image_path, \"rb\") as img_f:\n",
        "        image_data = img_f.read()\n",
        "    return image_data\n",
        "\n",
        "def top_k(predictions, map_proto_path, label_path, k=5):\n",
        "    # 创建节点ID --> 英文字符串查找。\n",
        "    node_lookup = tf_testing.NodeLookup(\n",
        "        label_lookup_path=map_proto_path, uid_lookup_path=label_path\n",
        "    )\n",
        "    # 打印 tensorflow 的 top5\n",
        "    top_k = predictions.argsort()[-k:][::-1]\n",
        "    # print(\"===== TENSORFLOW 结果 =======\")\n",
        "    return {\n",
        "        node_lookup.id_to_string(node_id): predictions[node_id]\n",
        "        for node_id in top_k\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "运行 TensorFlow 推理："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2023-06-09 14:12:21.225269: W tensorflow/core/framework/op_def_util.cc:369] Op BatchNormWithGlobalNormalization is deprecated. It will cease to work in GraphDef version 9. Use tf.nn.batch_normalization().\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /media/pc/data/lxw/ai/tvm/xinetzone/__pypackages__/3.10/lib/tvm/relay/testing/tf.py:136: convert_variables_to_constants (from tensorflow.python.framework.convert_to_constants) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate for instructions on how to migrate your code to TensorFlow v2.\n",
            "WARNING:tensorflow:From /media/pc/data/tmp/cache/conda/envs/tvmz/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py:952: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate for instructions on how to migrate your code to TensorFlow v2.\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2023-06-09 14:12:21.871263: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:353] MLIR V1 optimization pass is not enabled\n"
          ]
        }
      ],
      "source": [
        "with tf_compat_v1.Session() as sess:\n",
        "    # 从保存的 GraphDef 文件创建图\n",
        "    graph_def = create_graph(model_path)\n",
        "    # 为图中的节点添加形状属性。\n",
        "    graph_def = tf_testing.AddShapesToGraphDef(sess, \"softmax\")\n",
        "    softmax_tensor = sess.graph.get_tensor_by_name(\"softmax:0\")\n",
        "    image_data = read_image_tf(image_path)\n",
        "    predictions = sess.run(softmax_tensor, {\"DecodeJpeg/contents:0\": image_data})"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "展示结果："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "===== TENSORFLOW 结果 =======\n",
            "name                                              \tscore               \n",
            "-----------------------------------------------------------------\n",
            "African elephant, Loxodonta africana              \t0.58394\n",
            "tusker                                            \t0.33909\n",
            "Indian elephant, Elephas maximus                  \t0.03186\n",
            "banana                                            \t0.00022\n",
            "desk                                              \t0.00019\n"
          ]
        }
      ],
      "source": [
        "predictions = np.squeeze(predictions)\n",
        "results = top_k(predictions, map_proto_path, label_path, k=5)\n",
        "print(\"===== TENSORFLOW 结果 =======\")\n",
        "print(\"name\".ljust(50)+\"\\t\"+\"score\".ljust(20))\n",
        "print(\"-\"*65)\n",
        "for name, score in results.items():\n",
        "    print(f\"{name.ljust(50)}\\t{score:.5f}\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Relay 推理\n",
        "\n",
        "将 TensorFlow graph 定义导入到 Relay 前端。\n",
        "\n",
        "结果：\n",
        "\n",
        "- sym: Relay 表达式，表示给定的 tensorflow protobuf。\n",
        "- params: 从 tensorflow params（张量 protobuf）转换而来的参数。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "目标设备设置:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "import tvm\n",
        "from tvm import relay\n",
        "# 使用这些注释设置来构建 cuda\n",
        "# target = tvm.target.Target(\"cuda\", host=\"llvm\")\n",
        "# layout = \"NCHW\"\n",
        "# dev = tvm.cuda(0)\n",
        "target = tvm.target.Target(\"llvm\", host=\"llvm\")\n",
        "layout = None\n",
        "dev = tvm.cpu(0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/media/pc/data/lxw/ai/tvm/xinetzone/__pypackages__/3.10/lib/tvm/relay/frontend/tensorflow.py:537: UserWarning: Ignore the passed shape. Shape in graphdef will be used for operator DecodeJpeg/contents.\n",
            "  warnings.warn(\n",
            "/media/pc/data/lxw/ai/tvm/xinetzone/__pypackages__/3.10/lib/tvm/relay/frontend/tensorflow_ops.py:1036: UserWarning: DecodeJpeg: It's a pass through, please handle preprocessing before input\n",
            "  warnings.warn(\"DecodeJpeg: It's a pass through, please handle preprocessing before input\")\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "TensorFlow 的 protobuf 已导入到 Relay 前端。\n"
          ]
        }
      ],
      "source": [
        "shape = 299, 299, 3\n",
        "input_name = \"DecodeJpeg/contents\"\n",
        "shape_dict = {input_name: shape}\n",
        "dtype_dict = {input_name: \"uint8\"}\n",
        "with tf_compat_v1.Session() as sess:\n",
        "    # 从保存的 GraphDef 文件创建图\n",
        "    graph_def = create_graph(model_path)\n",
        "    # 为图中的节点添加形状属性。\n",
        "    graph_def = tf_testing.AddShapesToGraphDef(sess, \"softmax\")\n",
        "mod, params = relay.frontend.from_tensorflow(graph_def, layout=layout, shape=shape_dict)\n",
        "print(\"TensorFlow 的 protobuf 已导入到 Relay 前端。\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Relay 构建\n",
        "\n",
        "使用给定的输入规格将图编译为 LLVM 目标。 \n",
        "\n",
        "结果：\n",
        "  \n",
        "- `graph`：编译后的最终计算图。\n",
        "- `params`：编译后的最终参数。\n",
        "- `lib`：可以在具有 TVM 运行时的目标上部署的目标库。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with tvm.transform.PassContext(opt_level=3):\n",
        "    lib = relay.build(mod, target, params=params)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 在 TVM 上执行 portable graph\n",
        "\n",
        "现在我们可以尝试在目标设备上部署编译好的模型。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import graph_executor\n",
        "from PIL import Image\n",
        "\n",
        "image = Image.open(image_path).resize((299, 299))\n",
        "x = np.array(image)\n",
        "dtype = \"uint8\"\n",
        "m = graph_executor.GraphModule(lib[\"default\"](dev))\n",
        "# set inputs\n",
        "m.set_input(\"DecodeJpeg/contents\", tvm.nd.array(x.astype(dtype)))\n",
        "# execute\n",
        "m.run()\n",
        "# get outputs\n",
        "tvm_output = m.get_output(0)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### TVM 处理输出\n",
        "\n",
        "将 InceptionV1 模型的输出处理成可读的文本形式。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "===== TVM 结果 =======\n",
            "name                                              \tscore               \n",
            "-----------------------------------------------------------------\n",
            "African elephant, Loxodonta africana              \t0.58335\n",
            "tusker                                            \t0.33901\n",
            "Indian elephant, Elephas maximus                  \t0.02391\n",
            "banana                                            \t0.00025\n",
            "vault                                             \t0.00021\n"
          ]
        }
      ],
      "source": [
        "predictions = tvm_output.numpy()\n",
        "predictions = np.squeeze(predictions)\n",
        "results = top_k(predictions, map_proto_path, label_path, k=5)\n",
        "print(\"===== TVM 结果 =======\")\n",
        "print(\"name\".ljust(50)+\"\\t\"+\"score\".ljust(20))\n",
        "print(\"-\"*65)\n",
        "for name, score in results.items():\n",
        "    print(f\"{name.ljust(50)}\\t{score:.5f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 布局变换"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [],
      "source": [
        "# print(mod[\"main\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [],
      "source": [
        "desired_layouts = {\n",
        "    'image.resize2d': ['NCHW'],\n",
        "    'nn.conv2d': ['NCHW', 'default'],\n",
        "    'nn.max_pool2d': ['NCHW', 'default'],\n",
        "    'nn.avg_pool2d': ['NCHW', 'default'],\n",
        "}\n",
        "\n",
        "# 将布局转换为 NCHW\n",
        "# RemoveUnusedFunctions 用于清理图。\n",
        "seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),\n",
        "                                relay.transform.ConvertLayout(desired_layouts)])\n",
        "with tvm.transform.PassContext(opt_level=3):\n",
        "    mod = seq(mod)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [],
      "source": [
        "# print(mod[\"main\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Call relay compilation\n",
        "with relay.build_config(opt_level=3):\n",
        "     lib = relay.build(mod, target, params=params)\n",
        "m = graph_executor.GraphModule(lib[\"default\"](dev))\n",
        "# set inputs\n",
        "m.set_input(\"DecodeJpeg/contents\", tvm.nd.array(x.astype(dtype)))\n",
        "# execute\n",
        "m.run()\n",
        "# get outputs\n",
        "tvm_output = m.get_output(0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "===== TVM 结果 =======\n",
            "name                                              \tscore               \n",
            "-----------------------------------------------------------------\n",
            "African elephant, Loxodonta africana              \t0.58335\n",
            "tusker                                            \t0.33901\n",
            "Indian elephant, Elephas maximus                  \t0.02391\n",
            "banana                                            \t0.00025\n",
            "vault                                             \t0.00021\n"
          ]
        }
      ],
      "source": [
        "predictions = tvm_output.numpy()\n",
        "predictions = np.squeeze(predictions)\n",
        "results = top_k(predictions, map_proto_path, label_path, k=5)\n",
        "print(\"===== TVM 结果 =======\")\n",
        "print(\"name\".ljust(50)+\"\\t\"+\"score\".ljust(20))\n",
        "print(\"-\"*65)\n",
        "for name, score in results.items():\n",
        "    print(f\"{name.ljust(50)}\\t{score:.5f}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
